from typing import List

import numpy as np

from autograd import Module, Parameter, Tensor
from autograd.function import tanh
from autograd.optim import SGD

# ================================
# 🔢 数据编码函数
# ================================


def binary_encode(x: int) -> List[int]:
    """
    将整数 x 编码为 10 位二进制向量（最低位在前）
    例如: x = 5 → 二进制 101 → [1, 0, 1, 0, 0, 0, 0, 0, 0, 0]
    用途: 将数字转为神经网络可处理的输入（10维）
    """
    return [x >> i & 1 for i in range(10)]  # 取第 i 位（0~9）


def fizz_buzz_encode(x: int) -> List[int]:
    """
    将整数 x 编码为 4 维 one-hot 标签：
        [1, 0, 0, 0] → 普通数字（如 1, 2, 4...）
        [0, 1, 0, 0] → 能被 3 整除（fizz）
        [0, 0, 1, 0] → 能被 5 整除（buzz）
        [0, 0, 0, 1] → 能被 15 整除（fizzbuzz）
    """
    if x % 15 == 0:
        return [0, 0, 0, 1]  # fizzbuzz
    elif x % 5 == 0:
        return [0, 0, 1, 0]  # buzz
    elif x % 3 == 0:
        return [0, 1, 0, 0]  # fizz
    else:
        return [1, 0, 0, 0]  # 数字本身


# ================================
# 📊 构建训练数据集
# ================================

# 训练数据：101 ~ 1023（共 923 个样本）
# 输入：每个数字 → 10 维二进制向量
x_train = Tensor([binary_encode(x) for x in range(101, 1024)])
# x_train.shape = (923, 10)

# 标签：每个数字 → 4 维 one-hot 向量
y_train = Tensor([fizz_buzz_encode(x) for x in range(101, 1024)])
# y_train.shape = (923, 4)


# ================================
# 🧠 定义神经网络模型
# ================================


class FizzBuzzModule(Module):
    """
    一个简单的两层全连接神经网络：
        输入层 (10) → 隐藏层 (num_hidden) → 输出层 (4)
    使用 tanh 激活函数
    """

    def __init__(self, num_hidden: int = 50) -> None:
        # 第一层权重: (10, num_hidden)
        self.w1 = Parameter(10, num_hidden)  # 自动 requires_grad=True
        # 第一层偏置: (num_hidden,)
        self.b1 = Parameter(num_hidden)

        # 第二层权重: (num_hidden, 4)
        self.w2 = Parameter(num_hidden, 4)
        # 第二层偏置: (4,)
        self.b2 = Parameter(4)

    def forward(self, inputs: "Tensor") -> "Tensor":
        """
        前向传播
        inputs: (batch_size, 10)
        return: (batch_size, 4) —— 未归一化的 logits
        """
        # 线性变换1: (batch_size, 10) @ (10, num_hidden) → (batch_size, num_hidden)
        x1 = inputs @ self.w1 + self.b1
        # 激活函数: tanh → (batch_size, num_hidden)
        x2 = tanh(x1)
        # 线性变换2: (batch_size, num_hidden) @ (num_hidden, 4) → (batch_size, 4)
        x3 = x2 @ self.w2 + self.b2
        return x3

    def __call__(self, *args, **kwargs):
        """使 model(inputs) 等价于 model.forward(inputs)"""
        return self.forward(*args, **kwargs)


# ================================
# ⚙️ 训练配置
# ================================

model = FizzBuzzModule(num_hidden=50)  # 创建模型
optimizer = SGD(lr=0.001)  # 随机梯度下降优化器，学习率 0.001
epochs = 5000  # 训练轮数
batch_size = 32  # 每批样本数

# 生成批次起始索引: [0, 32, 64, ..., 928]（共 ceil(923/32)=29 批）
starts = np.arange(0, x_train.shape[0], batch_size)  # shape: (29,)


# ================================
# 🏃‍♂️ 开始训练
# ================================

for epoch in range(epochs):
    epoch_loss = 0.0

    # 每轮打乱批次顺序（增加随机性）
    np.random.shuffle(starts)

    for start in starts:
        end = start + batch_size

        # 清零模型中所有 Parameter 的梯度（避免累加）
        model.zero_grad()

        # 获取当前批次数据
        inputs = x_train[start:end]  # shape: (32, 10) 或更小（最后一批）
        actual = y_train[start:end]  # shape: (32, 4)

        # 前向传播：得到预测值
        predicted = model(inputs)  # shape: (32, 4)

        # 计算均方误差（MSE）损失
        errors = predicted - actual  # shape: (32, 4)
        loss = (errors * errors).sum()  # 标量（0维 Tensor）

        # 反向传播：自动计算所有 Parameter 的梯度
        loss.backward()

        # 累加损失（用于监控）
        epoch_loss += (
            loss.data.item() if hasattr(loss.data, "item") else float(loss.data)
        )

        # 更新参数：w = w - lr * grad
        optimizer.step(model)

    # 每轮打印损失
    print(f"Epoch {epoch + 1:4d} | Loss: {epoch_loss:.4f}")


# ================================
# 🧪 测试模型（1~100）
# ================================

print("\n" + "=" * 50)
print("🧪 测试结果（1 ~ 100）")
print("=" * 50)

num_correct = 0
labels = ["number", "fizz", "buzz", "fizzbuzz"]  # 对应 one-hot 索引

for x in range(1, 101):
    # 编码输入
    inputs = Tensor([binary_encode(x)])  # shape: (1, 10)

    # 模型预测
    predicted = model(inputs)  # shape: (1, 4)
    predicted_idx = np.argmax(predicted.data)  # 取最大值索引（0~3）

    # 真实标签
    actual_idx = np.argmax(fizz_buzz_encode(x))

    # 判断是否正确
    if predicted_idx == actual_idx:
        num_correct += 1
        status = "✅"
    else:
        status = "❌"

    print(
        f"{status} {x:2d}: 预测={labels[predicted_idx]:8s} | 真实={labels[actual_idx]}"
    )

print("\n" + "=" * 50)
print(f"🎯 正确率: {num_correct}/100 = {num_correct}%")
print("=" * 50)
